import cv2
import os
from preparation import *
from yolov5.utils.general import (non_max_suppression, scale_coords, cv2,
                                  xyxy2xywh)
from yolov5.utils.plots import Annotator, colors
from yolov5.utils.torch_utils import time_sync

def copeWithImg(path):
    img0 = cv2.imread(path)
    # img0 = cv2.resize(img0, (800, 600))
    # img0 = img0[230 * 5:500 * 5, :]
    # img0 = cv2.resize(img0, (800, 600))

    img = letterbox(img0, imgsz, stride=stride)[0]
    img = np.array([img])
    img = img[..., ::-1].transpose((0, 3, 1, 2))  # BGR to RGB, BHWC to BCHW
    img = np.ascontiguousarray(img)

    img = torch.from_numpy(img).to(device)

    img = img.half() if half else img.float()  # uint8 to fp16/32
    img /= 255.0

    pred = model(img, augment=augment, visualize=visualize)
    pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
    annotator = Annotator(img0, line_width=1, pil=not ascii, font_size=1)

    det = pred[0]

    if det is not None and len(det):
        det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.shape).round()

            # Print results
        for c in det[:, -1].unique():
            n = (det[:, -1] == c).sum()  # detections per class

        confs = det[:, 4]
        xyxys = det[:, 0:4]

        for j, (xyxy, conf) in enumerate(zip(xyxys, confs)):

            label = f'person {conf * 100:.2f}%'
            annotator.box_label(xyxy, label, color=colors(100, True))

    img0 = annotator.result()

    return img0, len(det)

def fileRead(path):
    svPthpre,name = path.rsplit('\\',maxsplit = 1)
    svFolder = svPthpre + 'label\\'
    if not os.path.exists(svFolder):
        os.mkdir(svFolder)

    img0,n = copeWithImg(path)

    path1 = svFolder+f'_{n}_'+name
    cv2.imshow('test',img0)
    cv2.waitKey(10)
    cv2.imwrite(path1,img0)
    print('copewith',path1)



def generate_labeled_picture(path):
    files = os.scandir(path)
    for fi in files:
        if fi.is_dir():
            generate_labeled_picture(fi.path)
        else:
            fileRead(fi.path)



generate_labeled_picture(r"E:\SMART_GLASS_WITH_DEPTHAI\person-recognize\picture\2")
